// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <kllm/utility/fs.h>
#include <cstdint>
#include <cstdio>
#include <string>
#include <algorithm>
#include <cinttypes>
#include <climits>
#include <cmath>
#include <codecvt>
#include <cstdarg>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <iterator>
#include <regex>
#include <sstream>
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <llama.h>

#if defined(__APPLE__) && defined(__MACH__)
#include <sys/types.h>
#include <sys/sysctl.h>
#endif

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#ifndef NOMINMAX
#   define NOMINMAX
#endif
#include <locale>
#include <windows.h>
#include <fcntl.h>
#include <io.h>
#else
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <unistd.h>
#endif
#if defined(LLAMA_USE_CURL)
#include <curl/curl.h>
#include <curl/easy.h>
#include <future>
#endif

#if defined(_MSC_VER)
#pragma warning(disable: 4244 4267) // possible loss of data
#endif

namespace kllm {


// Validate if a filename is safe to use
// To validate a full path, split the path by the OS-specific path separator, and validate each part with this function
    bool fs_validate_filename(const std::string & filename) {
        if (!filename.length()) {
            // Empty filename invalid
            return false;
        }
        if (filename.length() > 255) {
            // Limit at common largest possible filename on Linux filesystems
            // to avoid unnecessary further validation
            // (On systems with smaller limits it will be caught by the OS)
            return false;
        }

        std::u32string filename_utf32;
        try {
            std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
            filename_utf32 = converter.from_bytes(filename);

            // If the reverse conversion mismatches, it means overlong UTF-8 sequences were used,
            // or invalid encodings were encountered. Reject such attempts
            std::string filename_reencoded = converter.to_bytes(filename_utf32);
            if (filename_reencoded != filename) {
                return false;
            }
        } catch (const std::exception &) {
            return false;
        }

        // Check for forbidden codepoints:
        // - Control characters
        // - Unicode equivalents of illegal characters
        // - UTF-16 surrogate pairs
        // - UTF-8 replacement character
        // - Byte order mark (BOM)
        // - Illegal characters: / \ : * ? " < > |
        for (char32_t c : filename_utf32) {
            if (c <= 0x1F // Control characters (C0)
                || c == 0x7F // Control characters (DEL)
                || (c >= 0x80 && c <= 0x9F) // Control characters (C1)
                || c == 0xFF0E // Fullwidth Full Stop (period equivalent)
                || c == 0x2215 // Division Slash (forward slash equivalent)
                || c == 0x2216 // Set Minus (backslash equivalent)
                || (c >= 0xD800 && c <= 0xDFFF) // UTF-16 surrogate pairs
                || c == 0xFFFD // Replacement Character (UTF-8)
                || c == 0xFEFF // Byte Order Mark (BOM)
                || c == '/' || c == '\\' || c == ':' || c == '*' // Illegal characters
                || c == '?' || c == '"' || c == '<' || c == '>' || c == '|') {
                return false;
            }
        }

        // Reject any leading or trailing ' ', or any trailing '.', these are stripped on Windows and will cause a different filename
        // Unicode and other whitespace is not affected, only 0x20 space
        if (filename.front() == ' ' || filename.back() == ' ' || filename.back() == '.') {
            return false;
        }

        // Reject any ".." (currently stricter than necessary, it should be fine to just check for == ".." instead)
        if (filename.find("..") != std::string::npos) {
            return false;
        }

        // Reject "."
        if (filename == ".") {
            return false;
        }

        return true;
    }

// returns true if successful, false otherwise
    bool fs_create_directory_with_parents(const std::string & path) {
#ifdef _WIN32
        std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring wpath = converter.from_bytes(path);

    // if the path already exists, check whether it's a directory
    const DWORD attributes = GetFileAttributesW(wpath.c_str());
    if ((attributes != INVALID_FILE_ATTRIBUTES) && (attributes & FILE_ATTRIBUTE_DIRECTORY)) {
        return true;
    }

    size_t pos_slash = 0;

    // process path from front to back, procedurally creating directories
    while ((pos_slash = path.find('\\', pos_slash)) != std::string::npos) {
        const std::wstring subpath = wpath.substr(0, pos_slash);
        const wchar_t * test = subpath.c_str();

        const bool success = CreateDirectoryW(test, NULL);
        if (!success) {
            const DWORD error = GetLastError();

            // if the path already exists, ensure that it's a directory
            if (error == ERROR_ALREADY_EXISTS) {
                const DWORD attributes = GetFileAttributesW(subpath.c_str());
                if (attributes == INVALID_FILE_ATTRIBUTES || !(attributes & FILE_ATTRIBUTE_DIRECTORY)) {
                    return false;
                }
            } else {
                return false;
            }
        }

        pos_slash += 1;
    }

    return true;
#else
        // if the path already exists, check whether it's a directory
        struct stat info;
        if (stat(path.c_str(), &info) == 0) {
            return S_ISDIR(info.st_mode);
        }

        size_t pos_slash = 1; // skip leading slashes for directory creation

        // process path from front to back, procedurally creating directories
        while ((pos_slash = path.find('/', pos_slash)) != std::string::npos) {
            const std::string subpath = path.substr(0, pos_slash);
            struct stat info;

            // if the path already exists, ensure that it's a directory
            if (stat(subpath.c_str(), &info) == 0) {
                if (!S_ISDIR(info.st_mode)) {
                    return false;
                }
            } else {
                // create parent directories
                const int ret = mkdir(subpath.c_str(), 0755);
                if (ret != 0) {
                    return false;
                }
            }

            pos_slash += 1;
        }

        return true;
#endif // _WIN32
    }

    std::string fs_get_cache_directory() {
        std::string cache_directory = "";
        auto ensure_trailing_slash = [](std::string p) {
            // Make sure to add trailing slash
            if (p.back() != DIRECTORY_SEPARATOR) {
                p += DIRECTORY_SEPARATOR;
            }
            return p;
        };
        if (getenv("LLAMA_CACHE")) {
            cache_directory = std::getenv("LLAMA_CACHE");
        } else {
#ifdef __linux__
            if (std::getenv("XDG_CACHE_HOME")) {
                cache_directory = std::getenv("XDG_CACHE_HOME");
            } else {
                cache_directory = std::getenv("HOME") + std::string("/.cache/");
            }
#elif defined(__APPLE__)
            cache_directory = std::getenv("HOME") + std::string("/Library/Caches/");
#elif defined(_WIN32)
        cache_directory = std::getenv("LOCALAPPDATA");
#endif // __linux__
            cache_directory = ensure_trailing_slash(cache_directory);
            cache_directory += "llama.cpp";
        }
        return ensure_trailing_slash(cache_directory);
    }

    std::string fs_get_cache_file(const std::string & filename) {
        GGML_ASSERT(filename.find(DIRECTORY_SEPARATOR) == std::string::npos);
        std::string cache_directory = fs_get_cache_directory();
        const bool success = fs_create_directory_with_parents(cache_directory);
        if (!success) {
            throw std::runtime_error("failed to create cache directory: " + cache_directory);
        }
        return cache_directory + filename;
    }



}  // namespace kllm
